import os

from langchain_tavily import TavilySearch
from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder
from langchain.agents import create_tool_calling_agent,create_openai_tools_agent
from langchain.agents import AgentExecutor
from langchain_community.agent_toolkits import PlayWrightBrowserToolkit
from langchain_community.tools.playwright.utils import create_sync_playwright_browser
from langchain import hub



# 测试tavily搜索工具，tavily是一个聚合搜索引擎，集成了多种搜索引擎的结果，它将在网上搜索相关信息，确保环境变量存在TAVILY_API_KEY，默认免费账户一天可进行100次搜索
def test_tavily_search():
    search = TavilySearch(max_results=2)
    result = search.invoke("iPhone17有哪些功能")
    print("!!! Tavily Search Tool result:", result)

# 测试langchain agent中使用tavily搜索工具
def test_agent_with_tavily_tool():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    prompt =ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template("你是一名助人为乐的助手，并且可以调用工具进行网络搜索，获取实时信息。"),
        HumanMessagePromptTemplate.from_template("{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad")
    ])
    # 定义tavily工具
    search = TavilySearch(max_results=2)
    tools = [search]
    model.bind_tools(tools)
    # 创建langchain agent
    agent = create_tool_calling_agent(model, tools, prompt)
    # 使用agent executor来执行agent, verbose=True会打印中间过程
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    response = agent_executor.invoke({"input":"帮我搜索一下iPhone17有哪些功能？"})
    print(f"!!! Langchain agent with tavily tool, response:{response}, output:{response['output']}")




# playwright是一个浏览器自动化工具，可以用来抓取网页内容，模拟用户操作等
# playwright 需提前安装以下依赖 pip install playwright lxml langchain_community beautifulsoup4 reportlab
# 然后再执行 playwright install
def test_agent_with_playwright_tool():
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 创建playwright浏览器实例
    playwright_browser = create_sync_playwright_browser()
    # 定义工具
    toolkit = PlayWrightBrowserToolkit.from_browser(playwright_browser)
    tools = toolkit.get_tools()
    model.bind_tools(tools)
    # 定义prompt，prompt有社区上写好的，可以直接拉取使用
    prompt = hub.pull("hwchase17/openai-tools-agent")
    # 创建agent，支持OpenAI API RESTFUL API调用风格的模型，也可以使用create_openai_tools_agent
    agent = create_openai_tools_agent(model, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    # 执行任务
    response = agent_executor.invoke({"input":"请访问这个网站 https://www.apple.com/iphone-17/，总结一下iPhone17有哪些功能？"})
    print(f"!!! Langchain agent with playwright tool, response:{response}, output:{response['output']}")



if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    print(f"DEEPSEEK_API_KEY: {DEEPSEEK_API_KEY}")

    # test_tavily_search()
    test_agent_with_tavily_tool()
    test_agent_with_playwright_tool()







